import os
import subprocess
import time
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
from DisTools.Config import *

disassemblylogPath = logs["disLog"]
graphInfLog = logs["graphInfLog"]
logBasePath = logs["logPath"]
ida_log = logs["idaLog"]
errorLog = logs["errorLog"]
ida_script = script
samplesPath = data["samplesPath"]
savePath = data["dataSavePath"]
timeOut=maxProcessTime

class ida_disassembly_tool:

    def __init__(self,worknum=1) -> None:
        self.script=ida_script
        self.samplesPath=samplesPath
        self.workNum=worknum
        self.idaPath=idaPath
        self.scriptPath=ida_script
        self.space="      "
        self.savePath=savePath
        self.totalNum=0
        self.sampleList={}
        self.finishNum=0
        self.finishList={}
        self.unfinishedNum=0
        self.unfinishedList={}
        # 创建中间文件
        if not os.path.exists(savePath):
            os.mkdir(savePath)
        if not os.path.exists(logBasePath):
            os.mkdir(logBasePath)
        if not os.path.exists(disassemblylogPath):
            with open(disassemblylogPath,"w") as file:
                file.close()
        if not os.path.exists(graphInfLog):
            with open(graphInfLog,"w") as file:
                file.close()
        if not os.path.exists(ida_log):
            with open(ida_log,"w") as file:
                file.close()
        if not os.path.exists(errorLog):
            with open(errorLog,"w") as file:
                file.close()
        self.find_sample_to_disassembly()

    def analysisSample(self,samplePath,loc):
        """
            #运行如下命令
            #D:\software\IDA_Pro_7.7\ida64.exe -LE:\code\IDA_Pro\logs\ida_log.txt -A -SE:\code\IDA_Pro\ida_script.py \
            # D:\dataset\origin_customer_malware//976d3805e3a43064a686e4a10a00665f3db7316ef75a44ac7f553addeed9bff9
        """
        a=time.time()
        cmd = "{0} -L{1} -A -S{2} {3}".format(self.idaPath, ida_log, self.scriptPath, samplePath)
        print("执行命令"+cmd)
        subprocess.call(cmd)
        b=time.time()
        write_to_logs_file(disassemblylogPath,"{}{}处理完成****{}耗时: {}s\n".format(self.space,samplePath,self.space,b-a))

    def find_sample_to_disassembly(self):
        """
            找到需要处理的样本名，支持断点
            部分存在后缀，部分为md5
        """
        # suffixs = ["exe","dll","sys","DLL","EXE","SYS"]
        for file in  os.listdir(self.samplesPath):
            data=file.split(".")#找到样本文件
            if(len(data)==1):
                self.totalNum+=1
                self.sampleList[data[0]]=True
            # elif len(data)==2 and data[1] in suffixs:
            #     self.totalNum+=1
            #     self.sampleList[data[0]+"."+data[1]]=True
        for file in os.listdir(self.savePath):
            data=file.split(".")
            if(data[-1]=="bak"):
                if len(data)==2:
                    self.finishNum+=1
                    self.finishList[data[0]]=True
                # else:
                #     self.finishNum+=1
                #     self.finishList[data[0]+"."+data[1]]=True
        for name in self.sampleList.keys():
            if name not in self.finishList:
                self.unfinishedNum+=1
                self.unfinishedList[name]=self.samplesPath+"//"+name
        write_to_logs_file(disassemblylogPath,"样本总计{}个，已处理{}个,还剩{}个\n".format(self.totalNum,self.finishNum,self.unfinishedNum))  

    def run(self):
        thread_pool = ThreadPoolExecutor(max_workers=self.workNum)
        loc=0
        for name in self.unfinishedList.keys():
            loc+=1
            samplePath=self.unfinishedList[name]
            thread_pool.submit(self.analysisSample,samplePath,loc)
        thread_pool.shutdown(wait= True)  

def write_to_logs_file(disassemblylogPath,content):
    with open(disassemblylogPath,'a+',encoding="utf-8") as file:
        file.write(str(datetime.now())+str(content)+"\n")

if __name__ == '__main__':
    idadis=ida_disassembly_tool()
    idadis.run()